import os
import random
import json

rank_size = 8192
root = "/data01/huawei/zsz/resume_training/checkpoint"
for i in range(rank_size):
    print("process: rank ", i)
    rank_dir = os.path.join(root, f"rank_{i}")
    os.makedirs(rank_dir, exist_ok=True)
    meta_json = os.path.join(rank_dir, "meta.json")
    epoch = random.randint(900, 1000)
    step = 2
    ckpt = f"llama_7b_rank_{i}-{epoch}_{step}.ckpt"
    meta_data = {
        "last_epoch": epoch,
        "last_step": step,
        "last_ckpt_file": ckpt
    }

    with open(meta_json,"w") as f:
        json.dump(meta_data, f)

    for j in range(900, epoch+1):
        ckpt_path = os.path.join(rank_dir, f"llama_7b_rank_{i}-{j}_{step}.ckpt")
        f = open(ckpt_path, "w")
        f.close()